import torch
from torch import nn
import torchvision
import torchvision.datasets
import torchvision.transforms as transforms
import matplotlib.pyplot as plt
from torch.utils.data import Dataset, DataLoader
from torch.autograd import Variable
import torch.utils.data as Data
# Hyper Parameters
EPOCH = 5  # 训练整批数据多少次, 为了节约时间, 我们只训练一次
BATCH_SIZE = 100
TIME_STEP = 28  # rnn 时间步数 / 图片高度
INPUT_SIZE = 28  # rnn 每步输入值 / 图片每行像素
LR = 0.01  # learning rate
DOWNLOAD_MNIST = False  # 如果你已经下载好了mnist数据就写上 Fasle

# Mnist 手写数字
train_data = torchvision.datasets.MNIST(
    root='./mnist/',  # 保存或者提取位置
    train=True,  # this is training data
    transform=torchvision.transforms.ToTensor(),  # 转换 PIL.Image or numpy.ndarray 成
    # torch.FloatTensor (C x H x W), 训练的时候 normalize 成 [0.0, 1.0] 区间
    download=DOWNLOAD_MNIST
)

test_data = torchvision.datasets.MNIST(root='./mnist/', train=False)
# 批训练 50samples, 1 channel, 28x28 (50, 1, 28, 28)
train_loader = Data.DataLoader(dataset=train_data, batch_size=BATCH_SIZE, shuffle=True)
# 为了节约时间, 我们测试时只测试前2000个
test_x = torch.unsqueeze(test_data.data, dim=1).type(torch.FloatTensor)[:2000].cuda()/255.   # shape from (2000, 28, 28) to (2000, 1, 28, 28), value in range(0,1)
test_y = test_data.targets[:2000].cuda()


class RNN(nn.Module):
    def __init__(self):
        super(RNN, self).__init__()  # 继承RNN的init

        self.rnn = nn.LSTM(  # LSTM 效果要比 nn.RNN() 好多了
            input_size=28,  # 图片每行的数据像素点
            hidden_size=64,  # rnn hidden unit
            num_layers=1,  # 有几层 RNN layers   能力——计算时间
            batch_first=True,  # input & output 会是以 batch size 为第一维度的特征集 e.g.  (batch, time_step, input_size)
        )
        self.out = nn.Linear(64, 10)  # 输出层  10个数字分类

    def forward(self, x):
        # x shape (batch, time_step, input_size)
        # r_out shape (batch, time_step, output_size)
        # h_n shape (n_layers, batch, hidden_size)   LSTM 有两个 hidden states, h_n 是分线, h_c 是主线
        # h_c shape (n_layers, batch, hidden_size)
        r_out, (h_n, h_c) = self.rnn(x, None)  # None 表示 hidden state 会用全0的 state
        # 选取最后一个时间点的 r_out 输出
        # 这里 r_out[:, -1, :] 的值也是 h_n 的值
        out = self.out(r_out[:, -1, :])  # 选取最后一个output
        return out


rnn = RNN()
rnn.cuda()
optimizer = torch.optim.Adam(rnn.parameters(), lr=LR)  # optimize all parameters
loss_func = nn.CrossEntropyLoss()  # the target label is not one-hotted
for epoch in range(EPOCH):
    for step, (x, y) in enumerate(train_loader):   # gives batch data
        b_x = Variable(x.view(-1, 28, 28)) .cuda()  # reshape x to (batch, time_step, input_size)
        # print(b_x.size())
        b_y = Variable(y).cuda()
        output = rnn(b_x)               # rnn output
        # print(b_y.size())
        # print(output.size())
        loss = loss_func(output, b_y)   # cross entropy loss
        optimizer.zero_grad()           # clear gradients for this training step
        loss.backward()                 # backpropagation, compute gradients
        optimizer.step()                # apply gradients
        if step % 20 == 0:
            test_output = rnn(test_x.view(-1, 28, 28))  # 调整输入
            # print(test_x.size())
            # print(test_output.size())
            pred_y = torch.max(test_output, 1)[1].cuda().data.squeeze()
            # print(pred_y.size())
            right_sum = 0
            for a, b in zip(pred_y, test_y):
                if a == b:
                    right_sum += 1
                    accuracy = right_sum / test_y.size(0)
            print('Epoch:', epoch, '|step:', step, '|loss:', loss.data.cpu().numpy(), '|acc:', accuracy)

test_output = rnn(test_x[:10].view(-1, 28, 28))
pred_y = torch.max(test_output, 1)[1].data.cpu().numpy().squeeze()
print(pred_y, 'prediction number')
print(test_y[:10].cpu().numpy(), 'real number')